__doc__ =   """ Clustering of residues at different levels with different methods.
            """

from AbstractPropertyDict import AbstractAtomPropertyDict
from AbstractPropertyDict import AbstractResiduePropertyDict
from AbstractPropertyDict import AttributeDict
from Holders import AtomHolder
from Holders import ResidueHolder
from Einput import einput
from Selection import indicate
from Neighbors import Neighbors      # for k-nearest neighbors clustering
from scipy.cluster import vq
from numpy import mean
from functools import partial
from itertools import repeat, count, izip
from hcluster import *

def hcluster_call(coords, dmethod, dopts, dbool, lmethod, metric, **kwargs):
    """ This is a wrapper function for pdist, linkage, fcluster from hcluster.
    """
    if dopts:
        D = pdist(coords, dmethod, dopts)
    elif dbool:
        D = pdist(coords, dmethod)
    else:   # we must use euclidean distance and e.g ward on raw datapoints
        metric = 'euclidean'
        D = coords
    # it seems that for some reason a valid linkage cannot be build if the number
    # of coordinates is lower than 4. where should it be cought.
    # print 'done pdist %s %s' % (lmethod, metric)
    L = linkage(D, lmethod, metric)
    idx = fcluster(L, **kwargs)
    return idx

# partial wrappers around hcluster_call user has to provide coords and kwargs to fcluster
single = partial(hcluster_call, dmethod ='euclidean', dopts =None,\
                    dbool =True, lmethod ='single', metric ='euclidean')
average = partial(hcluster_call, dmethod ='euclidean', dopts =None,\
                    dbool =True, lmethod ='average', metric ='euclidean')
complete = partial(hcluster_call, dmethod ='euclidean', dopts =None,\
                    dbool =True, lmethod ='complete', metric ='euclidean')
ward = partial(hcluster_call, dmethod ='euclidean', dopts =None,\
                    dbool =False, lmethod ='ward', metric ='euclidean')
mahalanobis = partial(hcluster_call, dmethod ='mahalanobis', dopts =None,\
                    dbool =True, lmethod ='average', metric ='euclidean')

def cluster_h(entities, level, ctype, hmethod, **kwargs):
    """ Clusters entites using the hcluster library.
    """
    # partial functions provide a sane way to group settings.
    # you can create your own partial application by:
    # new_function = partial(hcluster_call, dmethod, dopts, dbool, lmethod, metric)
    entity_coords = AttributeDict(entities, ctype, level)   # entity must have the coords
    coords = entity_coords.array()                          # attribute already set.
    ids = entity_coords.keys()                              # key list from result dict
    idx = eval(hmethod)(coords, **kwargs)                   # give coords and kwargs
    cluster_map = dict(zip(ids, idx))                       # map results
    return cluster_map

def cluster_attr(entities, level, attr, **kwargs):
    """ Clusters entities using the _split_children interface by attribute value.
    """
    cluster_map = {}
    entity = einput(entities, level)
    split = entity._split_children(attr, **kwargs)  # entities split by value of attr
    for (cluster, num) in izip(split.itervalues(), count(1)): # enumerate from 1
        part = zip(cluster.keys(), repeat(num))               # assing properly
        cluster_map.update(part)
    return cluster_map

def cluster_kmeans(entities, level, ctype, k, **kwargs):
    """ Clusters entities using the numpy.vq.kmeans2 function.
    """
    entity_coords = AttributeDict(entities, ctype, level)       # entity must have the coords
    coords = entity_coords.array()                              # attribute already set.
    ids = entity_coords.keys()                                  # key list from result dict
    res, idx = vq.kmeans2(coords, k, **kwargs)                  # numpy
    cluster_map = dict(zip(ids, idx + 1))                       # map numpy result onto full_ids
    return cluster_map

def cluster_n(entities, level, ctype, points, **kwargs):
    """ Clusters entities by neighbors (proximity to other entities).
    """
    ents = einput(entities, level)
    cluster_map = {}
    knn = Neighbors(points, level, ctype)
    for (id_, entity) in ents.iteritems():
        nn = knn.get_neighbors(entity, 1, level)
        cluster_id = nn._data_children(**kwargs)[0]
        cluster_map.update({id_:cluster_id})
    return cluster_map


# Abstract Classes
class AbstractAtomClusters(AbstractAtomPropertyDict):
   def __init__(self, atoms, method, xtra_key, *args, **kwargs):
        """ Returns clusters of atoms, assumes atoms as input.
        """
        level = 'A'
        result = eval(method)(atoms, level, *args, **kwargs)    # atom id -> cluster id
        for (key, value) in result.iteritems():
            atoms[key].xtra[xtra_key] = value
        AbstractAtomPropertyDict.__init__(self, result)


class AbstractResidueClusters(AbstractResiduePropertyDict):
   def __init__(self, residues, method, xtra_key, *args, **kwargs):
        """ Returns clusters of residues, assumes residues as input.
        """
        level = 'R'
        result = eval(method)(residues, level, *args, **kwargs)    # res id -> cluster id
        for (key, value) in result.iteritems():
            residues[key].xtra[xtra_key] = value
        AbstractResiduePropertyDict.__init__(self, result)


class AbstractResidueFromAtomClusters(AbstractResiduePropertyDict):
    def __init__(self, entities, xtra_key):
        residues = einput(entities,'R', 'residues')
        cluster_map = {}
        for residue in residues:
            for atom in residue:
                cluster = atom.xtra.get(xtra_key)
                if cluster:
                    residue.xtra[xtra_key] = cluster
                    cluster_map.update({residue.get_full_id():cluster})
                    break
        AbstractResiduePropertyDict.__init__(self, cluster_map)


# Hierarchical depend on hcluster and use hcluster
class HierarchicalAtomClusters(AbstractAtomClusters):
    """ Hierarchically cluster atoms.
    """
    def __init__(self, entities, xtra_key ='H_CLUST', ctype ='coords', hmethod ='average', **kwargs):
        atoms = einput(entities, level='A', name ='atoms')
        kwargs['criterion'] = (kwargs.get('criterion')  or 'distance')
        kwargs['t']         = (kwargs.get('t')          or 15.0)
        method = 'cluster_h'
        AbstractAtomClusters.__init__(self, atoms, method, xtra_key, ctype, hmethod, **kwargs)


class HierarchicalResidueClusters(AbstractResidueClusters):
    """ Hierarchically cluster residues.
    """
    def __init__(self, entities, xtra_key ='H_CLUST', ctype ='coords', hmethod ='average', **kwargs):
        residues = einput(entities, level='R', name ='residues')
        kwargs['criterion'] = (kwargs.get('criterion')  or 'distance')
        kwargs['t']         = (kwargs.get('t')          or 15.0)
        method = 'cluster_h'
        AbstractResidueClusters.__init__(self, residues, method, xtra_key, ctype, hmethod, **kwargs)


# Kmeans depend on scipy
class KmeansAtomClusters(AbstractAtomClusters):
    """ Cluster atoms using the k-means algorithm.
    """
    def __init__(self, entities, xtra_key ='K_CLUST', ctype ='coords', k =8, **kwargs):
        atoms = einput(entities, level='A', name ='atoms')
        kwargs['minit'] = (kwargs.get('minit')  or 'points')
        method = 'cluster_kmeans'
        AbstractAtomClusters.__init__(self, atoms, method, xtra_key, ctype, k, **kwargs)


class KmeansResidueClusters(AbstractResidueClusters):
    """ Cluster residues using the k-means algorithm.
    """
    def __init__(self, entities, xtra_key ='K_CLUST', ctype ='coords', k =8, **kwargs):
        residues = einput(entities, level='R', name ='residues')
        kwargs['minit'] = (kwargs.get('minit')  or 'points')
        method = 'cluster_kmeans'
        AbstractResidueClusters.__init__(self, residues, method, xtra_key, ctype, k, **kwargs)


# Nearest Neighbor Clusters depend on ckd3 and Neighbors
class NearestAtomClusters(AbstractAtomClusters):
    def __init__(self, entities, points, xtra_key ='N_CLUST', ctype ='coords', **kwargs):
        atoms = einput(entities, level ='A', name ='atoms')
        kwargs['attr'] = (kwargs.get('attr') or 'ser_num')
        method = 'cluster_n'
        AbstractAtomClusters.__init__(self, atoms, method, xtra_key, ctype, points, **kwargs)


class NearestResidueClusters(AbstractResidueClusters):
    def __init__(self, entities, points, xtra_key ='N_CLUST', ctype ='coords', **kwargs):
        residues = einput(entities, level ='R', name ='residues')
        kwargs['attr'] = (kwargs.get('attr') or 'res_id')
        method = 'cluster_n'
        AbstractResidueClusters.__init__(self, residues, method, xtra_key, ctype, points, **kwargs)


# Attribute clusters are wrappers around _split_children
class AttrAtomClusters(AbstractAtomClusters):
    def __init__(self, entities, xtra_key ='A_CLUST', attr ='at_id', **kwargs):
        atoms = einput(entities, level ='A', name ='atoms')
        method = 'cluster_attr'
        AbstractAtomClusters.__init__(self, atoms, method, xtra_key, attr, **kwargs)


class AttrResidueClusters(AbstractResidueClusters):
    def __init__(self, entities, xtra_key ='A_CLUST', attr ='res_id', **kwargs):
        residues = einput(entities, level ='R', name ='residues')
        method = 'cluster_attr'
        AbstractResidueClusters.__init__(self, residues, method, xtra_key, attr, **kwargs)


#EOF